ffe283
@@ -1,7 +1,22 @@
+/*
+ * Copyright 2002-2012 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
 package org.springframework.security.web.context;
 
 import java.io.IOException;
+import java.io.PrintWriter;
+import java.io.Writer;
 
+import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpServletResponseWrapper;
 
@@ -10,8 +25,9 @@
import org.springframework.security.core.context.SecurityContextHolder;
 
 /**
  * Base class for response wrappers which encapsulate the logic for storing a security context and which
- * store the with the <code>SecurityContext</code> when a <code>sendError()</code> or <code>sendRedirect</code>
- * happens. See issue SEC-398.
+ * store the <code>SecurityContext</code> when a <code>sendError()</code>, <code>sendRedirect</code>,
+ * <code>getOutputStream().close()</code>, <code>getOutputStream().flush()</code>, <code>getWriter().close()</code>, or
+ * <code>getWriter().flush()</code> happens. See issue SEC-398 and SEC-2005.
  * <p>
  * Sub-classes should implement the {@link #saveContext(SecurityContext context)} method.
  * <p>
@@ -19,6 +35,7 @@
import org.springframework.security.core.context.SecurityContextHolder;
  *
  * @author Luke Taylor
  * @author Marten Algesten
+ * @author Rob Winch
  * @since 3.0
  */
 public abstract class SaveContextOnUpdateOrErrorResponseWrapper extends HttpServletResponseWrapper {
@@ -74,6 +91,34 @@
public abstract class SaveContextOnUpdateOrErrorResponseWrapper extends HttpServ
         super.sendRedirect(location);
     }
 
+    /**
+     * Makes sure the context is stored before calling <code>getOutputStream().close()</code> or
+     * <code>getOutputStream().flush()</code>
+     */
+    @Override
+    public ServletOutputStream getOutputStream() throws IOException {
+        return new SaveContextServletOutputStream(super.getOutputStream());
+    }
+
+    /**
+     * Makes sure the context is stored before calling <code>getWriter().close()</code> or
+     * <code>getWriter().flush()</code>
+     */
+    @Override
+    public PrintWriter getWriter() throws IOException {
+        return new SaveContextPrintWriter(super.getWriter());
+    }
+
+    /**
+     * Makes sure the context is stored before calling the
+     * superclass <code>flushBuffer()</code>
+     */
+    @Override
+    public void flushBuffer() throws IOException {
+        doSaveContext();
+        super.flushBuffer();
+    }
+
     /**
      * Calls <code>saveContext()</code> with the current contents of the <tt>SecurityContextHolder</tt>.
      */
@@ -115,10 +160,59 @@
public abstract class SaveContextOnUpdateOrErrorResponseWrapper extends HttpServ
     }
 
     /**
-     * Tells if the response wrapper has called <code>saveContext()</code> because of an error or redirect.
+     * Tells if the response wrapper has called <code>saveContext()</code> because of this wrapper.
      */
     public final boolean isContextSaved() {
         return contextSaved;
     }
 
+    /**
+     * Ensures the {@link SecurityContext} is updated prior to methods that commit the response.
+     * @author Rob Winch
+     */
+    private class SaveContextPrintWriter extends PrintWriter {
+
+        public SaveContextPrintWriter(Writer out) {
+            super(out);
+        }
+
+        public void flush() {
+            doSaveContext();
+            super.flush();
+        }
+
+        public void close() {
+            doSaveContext();
+            super.close();
+        }
+    }
+
+    /**
+     * Ensures the {@link SecurityContext} is updated prior to methods that commit the response.
+     *
+     * @author Rob Winch
+     */
+    private class SaveContextServletOutputStream extends ServletOutputStream {
+        private final ServletOutputStream delegate;
+
+        public SaveContextServletOutputStream(ServletOutputStream delegate) {
+            this.delegate = delegate;
+        }
+
+        public void write(int b) throws IOException {
+            this.delegate.write(b);
+        }
+
+        @Override
+        public void flush() throws IOException {
+            doSaveContext();
+            super.flush();
+        }
+
+        @Override
+        public void close() throws IOException {
+            doSaveContext();
+            super.close();
+        }
+    }
 }
